import cv2
from ultralytics import YOLO
from collections import defaultdict


def box_label(image, box, label='', color=(128, 128, 128), txt_color=(255, 255, 255)):
    # 得到目标矩形框的左上角和右下角坐标
    p1, p2 = (int(box[0]), int(box[1])), (int(box[2]), int(box[3]))
    # 绘制矩形框
    cv2.rectangle(image, p1, p2, color, thickness=1, lineType=cv2.LINE_AA)
    if label:
        # 得到要书写的文本的宽和长，用于给文本绘制背景色
        w, h = cv2.getTextSize(label, 0, fontScale=2 / 3, thickness=1)[0]
        # 确保显示的文本不会超出图片范围
        outside = p1[1] - h >= 3
        p2 = p1[0] + w, p1[1] - h - 3 if outside else p1[1] + h + 3
        cv2.rectangle(image, p1, p2, color, -1, cv2.LINE_AA)  # 填充颜色
        # 书写文本
        cv2.putText(image,
                    label, (p1[0], p1[1] - 2 if outside else p1[1] + h + 2),
                    0,
                    2 / 3,
                    txt_color,
                    thickness=1,
                    lineType=cv2.LINE_AA)


if __name__ == '__main__':
    # track_history用于保存目标ID，以及它在各帧的目标位置坐标，这些坐标是按先后顺序存储的
    track_history = defaultdict(lambda: [])

    model = YOLO('best.pt')
    # cap = cv2.VideoCapture(R"G:\temp\tank\3.mp4")
    # cap = cv2.VideoCapture(R"F:\Work\other\LiChao\original\方位90俯仰45\方位90俯仰45_gray0.mp4")
    cap = cv2.VideoCapture(1)

    fps = cap.get(cv2.CAP_PROP_FPS)
    size = (int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)), int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)))
    fNUMS = cap.get(cv2.CAP_PROP_FRAME_COUNT)

    # 视频帧循环
    while cap.isOpened():
        # 读取一帧图像
        success, frame = cap.read()

        if success:
            # 在帧上运行YOLOv8跟踪，persist为True表示保留跟踪信息，conf为0.3表示只检测置信值大于0.3的目标
            results = model.track(frame, conf=0.25, persist=True)
            if len(results[0]) == 0 or results[0].boxes.id == None:
                cv2.imshow("YOLOv8 Tracking", frame)  # 显示标记好的当前帧图像
                cv2.waitKey(15)
                continue
            # 得到该帧的各个目标的ID
            track_ids = results[0].boxes.id.int().cpu().tolist()
            # 遍历该帧的所有目标
            for track_id, box in zip(track_ids, results[0].boxes.data):
                if box[-1] == 80:  # 目标坦克
                    # 绘制该目标的矩形框
                    box_label(frame, box, '#' + str(track_id) + ' tank', (167, 146, 11))
                    # 得到该目标矩形框的中心点坐标(x, y)
                    x1, y1, x2, y2 = box[:4]
                    x = (x1 + x2) / 2
                    y = (y1 + y2) / 2
                    # 提取出该ID的以前所有帧的目标坐标，当该ID是第一次出现时，则创建该ID的字典
                    track = track_history[track_id]
                    track.append((float(x), float(y)))  # 追加当前目标ID的坐标

                elif box[-1] == 0:  # 目标为人
                    box_label(frame, box, '#' + str(track_id) + ' person', (67, 161, 255))

                    x1, y1, x2, y2 = box[:4]
                    x = (x1 + x2) / 2
                    y = (y1 + y2) / 2
                    track = track_history[track_id]
                    track.append((float(x), float(y)))  # x, y center point

            cv2.imshow("YOLOv8 Tracking", frame)  # 显示标记好的当前帧图像

            if cv2.waitKey(5) & 0xFF == ord("q"):  # 'q'按下时，终止运行
                break

        else:  # 视频播放结束时退出循环
            break

    # 释放视频捕捉对象，并关闭显示窗口
    cap.release()
    cv2.destroyAllWindows()
